import sys
import os
import cv2
import numpy as np
from PyQt5.QtWidgets import QApplication, QWidget, QHBoxLayout,QVBoxLayout, QPushButton, QLabel, QFileDialog,QLineEdit,QMessageBox
from PyQt5.QtGui import QPixmap
from paddleocr import PaddleOCR
from datetime import datetime

crop_img = 'crop_img'

class ImageTextExtractor(QWidget):
    def __init__(self):
        super().__init__()

        self.init_ui()
        self.image_folder_path = ""
        # self.ocr = PaddleOCR(use_gpu=True, rec_model_dir="D:/Data/rec_ppocr_v4/inference")
        self.ocr = PaddleOCR(use_gpu=True)

    def init_ui(self):
        # 创建垂直布局管理器
        layout = QVBoxLayout()

         # 创建用于放置选择文件夹路径按钮和开始按钮的水平布局管理器
        button_layout = QHBoxLayout()

        # 创建一个按钮，用于选择图片所在文件夹
        self.select_folder_button = QPushButton("选择图片所在文件夹")
        self.select_folder_button.clicked.connect(self.select_image_folder)

        # 创建一个标签，用于显示所选的文件夹路径
        self.folder_path_label = QLabel("未选择文件夹")

        # 创建"开始"按钮
        self.start_button = QPushButton("开始")
        self.start_button.clicked.connect(self.start_extraction)

        # 将选择文件夹路径按钮和开始按钮添加到水平布局中
        button_layout.addWidget(self.select_folder_button)
        button_layout.addWidget(self.start_button)


        # 创建一个文本输入框，用于输入图片所在文件夹路径
        self.path_input_box = QLineEdit()
        # 创建"显示"按钮
        self.show_button = QPushButton("显示")
        self.show_button.clicked.connect(self.show_images)

        button_layout.addWidget(self.show_button)
        self.image_label = QLabel()

        # 创建一个QWidget容器来放置水平布局
        button_widget = QWidget()
        button_widget.setLayout(button_layout)

        # 将按钮和标签添加到垂直布局中
        layout.addWidget(self.folder_path_label)
        layout.addWidget(button_widget)
        # layout.addWidget(self.select_folder_button)
        # layout.addWidget(self.start_button)
        layout.addWidget(self.path_input_box)
        # layout.addWidget(self.show_button)
        layout.addWidget(self.image_label)

        # 设置窗口的布局
        self.setLayout(layout)

        # 设置窗口标题和大小
        self.setWindowTitle("图片文字提取工具")
        # self.setFixedSize(480, 200)
        self.setMinimumSize(1444, 480)  # 设置最小尺寸，可根据需要调整

        # 使窗口可调节大小
        self.setSizePolicy(1, 1)

    def select_image_folder(self):
        # 打开文件对话框，选择文件夹
        folder_path = QFileDialog.getExistingDirectory(self, "选择图片所在文件夹")

        if folder_path:
            self.image_folder_path = folder_path
            self.folder_path_label.setText(folder_path)

    def start_extraction(self):
        if not self.image_folder_path:
            print("请先选择图片所在文件夹！")
            return

        # 获取文件夹中所有图片文件的路径
        image_paths = [os.path.join(self.image_folder_path, f) for f in os.listdir(self.image_folder_path) if f.endswith(('.jpg', '.png', '.jpeg'))]

        # 按图片名升序排序
        image_paths.sort()

        # 获取当前时间并格式化为字符串，作为文档名
        current_time = datetime.now().strftime("%Y%m%d%H%M")
        output_doc_path =  self.image_folder_path + f'_gt_{current_time}.txt'
        train_rec_path = self.image_folder_path + f'_crop_{current_time}.txt'
        with open(output_doc_path, 'w', encoding='utf-8') as output_doc:
            with open(train_rec_path, 'w', encoding='utf-8') as train_rec:
                for image_path in image_paths:
                    try:
                        img = cv2.imread(image_path)
                        # 使用PaddleOCR进行文字识别
                        result = self.ocr.ocr(img)
                        save_dir = self.image_folder_path + '/' + crop_img + '/'
                        if not os.path.exists(save_dir):
                            os.makedirs(save_dir)
                        # 提取识别出的文字内容
                        text = ""
                        index = 0
                        for line in result[0]:
                            index += 1
                            text += line[1][0] + "\n"
                            print(line[0], ":", line[1][0])
                            # polygon_vertices = np.array(line[0], dtype=np.int32)
                            # minx = int(min(line[0][0][0], line[0][3][0]))
                            # maxx = int(max(line[0][1][0], line[0][2][0]))
                            # miny = int(min(line[0][0][1], line[0][1][1]))
                            # maxy = int(max(line[0][2][1], line[0][3][1]))
                            # crop_img = img[miny:maxy, minx:maxx]
                            polygon_vertices = np.array(line[0], dtype=np.int32)
                            mask = np.zeros(img.shape[:2], dtype=np.uint8)
                            cv2.fillPoly(mask, [polygon_vertices], 255)
                            # print("== ", line[0][0][0], line[0][0][1],line[0][0][1],line[0][1][1],line[0][2][1],line[0][3][1])
                            cropped_image = cv2.bitwise_and(img, img, mask=mask)
                            # 获取多边形区域的最小外接矩形
                            x, y, w, h = cv2.boundingRect(np.array(polygon_vertices, dtype=np.int32))

                            # 从裁剪后的图像中提取出最小外接矩形区域内的图像，即只保留需要的部分
                            final_cropped_image = cropped_image[y:y + h, x:x + w]

                            # crop_img = img[line[0][0][0]:line[0][1][0], min(line[0][0][1],line[0][1][1])
                                        #    :max(line[0][2][1],line[0][3][1])]
                            # crop_img_path = os.path.pardir(image_path).join("crop_img").join(os.path.basename(image_path) + "{:d}".format(index) + os.path.splitext(image_path)[1])
                            crop_relative_path = crop_img + "/" + os.path.splitext(os.path.basename(image_path))[0] + "_" + str(index) + os.path.splitext(os.path.basename(image_path))[1]

                            crop_img_path = os.path.dirname(image_path) + "/" + crop_relative_path
                            print("save to ", crop_img_path)
                            cv2.imwrite(crop_img_path, final_cropped_image)
                            train_rec.write(crop_relative_path + '\t' + line[1][0] + '\n')
                        # print(line[0],text)
                        # 将图片路径和提取的文字写入文档
                        output_doc.write(f"图片路径: {image_path}\n")
                        output_doc.write(f"提取文字:\n{text}\n\n")
                    except Exception as e:
                        print(f"处理图片 {image_path} 时出错: {e}")

    def show_images(self):
        if not self.image_folder_path:
            QMessageBox.information(self, "选择文件夹", "请选择文件夹")
            return
        image_path = self.path_input_box.text()
        if not image_path:
            print("请输入有效的文件夹路径！")
            QMessageBox.information(self, "选择文件夹", "请输入文件路径")
            return

        # image_files = [f for f in os.listdir(folder_path) if f.endswith(('.jpg', '.png', '.jpeg'))]
        # if not image_files:
        #     print("该文件夹下没有符合要求的图片！")
        #     return

        show_image_path = os.path.join(self.image_folder_path, image_path)
        if not os.path.exists(show_image_path):
            QMessageBox.information(self, "文件不存在",f'请查看：{show_image_path} 是否存在')
            return
        pixmap = QPixmap(show_image_path)
        self.image_label.setPixmap(pixmap)
        self.image_label.setScaledContents(True)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = ImageTextExtractor()
    window.show()
    sys.exit(app.exec_())
